import configparser
import multiprocessing
from datetime import datetime, timedelta
from joblib import Parallel, delayed
from tqdm import tqdm
from stock_a.common.db import stock_info
from stock_a.feature.global_feature_source import GlobalFeatureSource
from stock_a.feature.stock_feature_source import StockFeatureSource
from stock_a.feature import feature_pb2
from stock_a.feature import feature_resigteror


class StockFeatureGenerator:
    def __init__(self):
        self.quant_config = configparser.ConfigParser()
        self.quant_config.read('../config/quant_config.ini')
        self.feature_generator_class_list = feature_resigteror.get_feature_generator_class_list()

    def init_tables(self):
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        for i in range(20):
            cursor.execute(stock_feature_create_sql.format(i))
        db.close()

    def clear_all_data(self):
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        for i in range(20):
            cursor.execute('DROP TABLE IF EXISTS stock_feature_day_{}'.format(str(i)))
            db.commit()
        cursor.close()
        db.close()

    def build_or_update_stock_feature(self, biz_date):
        global_feature_source = GlobalFeatureSource()
        # self.__execute_update_stock_feature('600900', global_feature_source, biz_date)
        code_set = stock_info.get_all_stock_code_set()
        Parallel(n_jobs=multiprocessing.cpu_count())(delayed(StockFeatureGenerator.__execute_update_stock_feature)(
            self, code, global_feature_source, biz_date) for code in tqdm(code_set, '构建/更新股票样本特征'))

    def __execute_update_stock_feature(self, code, global_feature_source, biz_date):
        stock_feature_source = StockFeatureSource(code)
        earliest_date = self.quant_config['stock_info']['earliest_date']
        table_name = stock_info.get_stock_feature_day_table_name(code)
        db = stock_info.new_db_conn()
        cursor = db.cursor()
        cursor.execute('SELECT `date` FROM {} WHERE `code` = %s ORDER BY date DESC LIMIT 1'.
                       format(table_name), code)
        data = cursor.fetchone()
        start_date = (data[0] + timedelta(days=1)).strftime('%Y%m%d') if data is not None else datetime.strptime(earliest_date, '%Y%m%d').date()
        total_num = len(stock_feature_source.info_group_idx_list)
        feature_generators = list()
        for feature_generator_class in self.feature_generator_class_list:
            feature_generators.append(feature_generator_class(stock_feature_source, global_feature_source))
        feature_db_data_list = list()
        for i in range(total_num):
            date = stock_feature_source.info_group_idx_list[i][0]
            if date < start_date:
                continue
            if date > biz_date:
                break
            feature_pb = feature_pb2.feature()
            for feature_generator in feature_generators:
                feature_generator.set_pb_feature(i, feature_pb)
            feature_binary = feature_pb.SerializeToString()
            feature_db_data_list.append([code, date, feature_binary])
        if len(feature_db_data_list) > 0:
            sql = (('INSERT INTO {} (`code`, `date`, `feature`) '
                    'VALUES (%s, %s, %s)'
                    ).format(
                table_name))
            cursor.executemany(sql, feature_db_data_list)
            db.commit()
        cursor.close()
        db.close()


stock_feature_create_sql = '''
create table if not exists `stock_feature_day_{}` (
    `code` VARCHAR(20) NOT NULL COMMENT '股票代码',
    `date` TIMESTAMP NOT NULL COMMENT '样本日期',
    `feature` BLOB NOT NULL COMMENT '样本特征（protobuf编码）',
    KEY (`code`),
    PRIMARY KEY (`code`, `date`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8 COMMENT '股票样本特征流水表';
'''